import math
from collections import Counter
from operator import itemgetter

class KNN:
    """
    K近邻算法实现（仅使用Python标准库）
    适用于分类问题
    """
    
    def __init__(self, k=5, weights='uniform', metric='euclidean'):
        """
        初始化KNN分类器
        
        参数:
        k: 邻居数量，默认5
        weights: 权重类型，'uniform'（平等）或'distance'（按距离加权）
        metric: 距离度量，'euclidean'（欧氏距离）或'manhattan'（曼哈顿距离）
        """
        self.k = k
        self.weights = weights
        self.metric = metric
        self.X_train = None
        self.y_train = None
    
    def fit(self, X, y):
        """
        训练模型（KNN只需存储数据）
        
        参数:
        X: 训练特征，列表的列表或元组的列表
        y: 训练标签，列表
        """
        self.X_train = [tuple(sample) for sample in X]  # 转换为元组确保不可变
        self.y_train = list(y)
        return self
    
    def _calculate_distance(self, x1, x2):
        """
        计算两个样本点之间的距离[1,3](@ref)
        """
        if len(x1) != len(x2):
            raise ValueError("样本特征维度不一致")
        
        if self.metric == 'euclidean':
            # 欧氏距离: sqrt(Σ(x_i - y_i)^2)[1](@ref)
            squared_diff = sum((a - b) ** 2 for a, b in zip(x1, x2))
            return math.sqrt(squared_diff)
        elif self.metric == 'manhattan':
            # 曼哈顿距离: Σ|x_i - y_i|[1](@ref)
            return sum(abs(a - b) for a, b in zip(x1, x2))
        else:
            raise ValueError("不支持的距离度量方法")
    
    def _normalize_minmax(self, X):
        """
        最小-最大归一化[2,7](@ref)
        公式: (x - min) / (max - min)
        """
        if not X:
            return X
        
        # 转置矩阵以便按特征列计算
        transposed = list(zip(*X))
        normalized_transposed = []
        
        for feature_col in transposed:
            min_val = min(feature_col)
            max_val = max(feature_col)
            range_val = max_val - min_val
            
            if range_val == 0:  # 避免除零
                normalized_col = [0.0] * len(feature_col)
            else:
                normalized_col = [(x - min_val) / range_val for x in feature_col]
            
            normalized_transposed.append(normalized_col)
        
        # 转置回来
        return list(zip(*normalized_transposed))
    
    def predict_single(self, x):
        """
        预测单个样本的类别[1,3,8](@ref)
        """
        if self.X_train is None or self.y_train is None:
            raise ValueError("请先使用fit方法训练模型")
        
        # 计算与所有训练样本的距离
        distances = []
        for i, train_sample in enumerate(self.X_train):
            dist = self._calculate_distance(x, train_sample)
            distances.append((dist, i))
        
        # 按距离排序并选择前k个[3](@ref)
        distances.sort(key=itemgetter(0))
        k_nearest = distances[:self.k]
        
        # 统计k个邻居的类别[1](@ref)
        if self.weights == 'uniform':
            # 平等权重
            class_votes = {}
            for dist, idx in k_nearest:
                label = self.y_train[idx]
                class_votes[label] = class_votes.get(label, 0) + 1
            
            # 返回出现次数最多的类别[3](@ref)
            return max(class_votes.items(), key=itemgetter(1))[0]
        
        elif self.weights == 'distance':
            # 距离加权（距离越小权重越大）
            class_weights = {}
            for dist, idx in k_nearest:
                label = self.y_train[idx]
                # 避免除零，给距离加一个小值
                weight = 1 / (dist + 1e-8)
                class_weights[label] = class_weights.get(label, 0) + weight
            
            return max(class_weights.items(), key=itemgetter(1))[0]
    
    def predict(self, X):
        """
        预测多个样本的类别
        """
        return [self.predict_single(tuple(sample)) for sample in X]
    
    def score(self, X_test, y_test):
        """
        计算模型在测试集上的准确率
        """
        predictions = self.predict(X_test)
        correct = sum(1 for pred, true in zip(predictions, y_test) if pred == true)
        return correct / len(y_test)

# 示例使用
if __name__ == "__main__":
    # 创建示例数据集[2,5](@ref)
    X_train = [
        [1.0, 1.1],
        [1.0, 1.0], 
        [0.0, 0.0],
        [0.0, 0.1]
    ]
    y_train = ['A', 'A', 'B', 'B']
    
    # 创建KNN分类器
    knn = KNN(k=3, weights='uniform', metric='euclidean')
    
    # 训练模型
    knn.fit(X_train, y_train)
    
    # 预测新样本
    test_sample = [0.2, 0.2]
    prediction = knn.predict_single(test_sample)
    print(f"样本 {test_sample} 的预测类别: {prediction}")
    
    # 测试多个样本
    test_samples = [[0.5, 0.5], [1.2, 1.0]]
    predictions = knn.predict(test_samples)
    print(f"批量预测结果: {predictions}")
    
    # 测试准确率
    X_test = [[0.1, 0.1], [1.1, 1.0]]
    y_test = ['B', 'A']
    accuracy = knn.score(X_test, y_test)
    print(f"模型准确率: {accuracy:.2%}")